import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim
import torchvision.transforms as transforms
from torchsummary import summary
from data import load_cifar10_data
from train import train_one_epoch
import matplotlib.pyplot as plt
from PIL import Image


class Net(nn.Module):
    def __init__(self, num_classes, in_channels):
        super(Net, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(8192, 1024),
            nn.Dropout(),
            nn.Linear(1024, num_classes)
        )

    def forward(self, input):
        y = self.conv1(input)
        y = self.conv2(y)
        y = self.classifier(y)
        return y
    
IMG_SIZE, IMG_CHANNELS = 32, 3
BATCH_SIZE = 16
train_loader, val_loader, _, classes = load_cifar10_data(BATCH_SIZE, IMG_SIZE)


model = Net(10, 3)
summary(model, (IMG_CHANNELS, IMG_SIZE, IMG_SIZE), BATCH_SIZE)


img = Image.open('./images/dog.jpeg')
resized_img = img.resize((IMG_SIZE, IMG_SIZE))
x = torch.unsqueeze(torch.from_numpy(np.array(resized_img)).permute(2, 0, 1).float()/255.0, dim=0)
transform = transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
x = transform(x)


plt.ion()
plt.figure(figsize=(10, 10))
plt.axis('off')

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.StepLR(optimizer, 2, 0.9)

for epoch in range(10):

    model.to('cpu')

    model.eval()
    y = model.conv1(x)
    y = model.conv2(y).detach()
    y = torch.squeeze(y)

    plt.clf()
    for i in range(y.shape[0]):
        fig = plt.subplot(8, 16, i + 1)
        fig.imshow(y[i], cmap='gray')
        fig.axis('off')
    plt.pause(0.1)

    y_hat = model(x).argmax()
    print(y_hat, classes[y_hat])

    train_one_epoch(epoch, model, train_loader, val_loader, criterion, optimizer, torch.device('mps'))

    scheduler.step()
    print(', lr={:.4f}'.format(scheduler.get_last_lr()[0]))

plt.pause(10)